#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
 
/*
 * 多项式拟合函数，等效于numpy.polyfit
 * x: 输入x值数组
 * y: 输入y值数组
 * n: 数据点数量
 * deg: 多项式阶数
 * coef: 输出的多项式系数数组，长度为deg+1
 * 返回值: 0表示成功，非0表示失败
 */
int polyfit(const double* x, const double* y, int n, int deg, double* coef) {
    if (n <= deg) {
        fprintf(stderr, "错误: 数据点数量必须大于多项式阶数\n");
        return 1;
    }

    // 分配内存用于正规方程矩阵和右侧向量
    double* A = (double*)malloc((deg + 1) * (deg + 1) * sizeof(double));
    double* b = (double*)malloc((deg + 1) * sizeof(double));
    if (!A || !b) {
        fprintf(stderr, "内存分配失败\n");
        free(A);
        free(b);
        return 1;
    }

    // 初始化正规方程矩阵和右侧向量
    memset(A, 0, (deg + 1) * (deg + 1) * sizeof(double));
    memset(b, 0, (deg + 1) * sizeof(double));

    // 填充正规方程矩阵和右侧向量
    for (int i = 0; i <= deg; i++) {
        for (int j = 0; j <= deg; j++) {
            double sum = 0.0;
            for (int k = 0; k < n; k++) {
                sum += pow(x[k], i + j);
            }
            A[i * (deg + 1) + j] = sum;
        }

        double sum = 0.0;
        for (int k = 0; k < n; k++) {
            sum += pow(x[k], i) * y[k];
        }
        b[i] = sum;
    }

    // 使用高斯消元法求解正规方程 Ax = b
    for (int i = 0; i <= deg; i++) {
        // 寻找主元
        double max_val = fabs(A[i * (deg + 1) + i]);
        int max_row = i;
        for (int j = i + 1; j <= deg; j++) {
            if (fabs(A[j * (deg + 1) + i]) > max_val) {
                max_val = fabs(A[j * (deg + 1) + i]);
                max_row = j;
            }
        }

        // 交换行
        if (max_row != i) {
            for (int j = i; j <= deg; j++) {
                double temp = A[i * (deg + 1) + j];
                A[i * (deg + 1) + j] = A[max_row * (deg + 1) + j];
                A[max_row * (deg + 1) + j] = temp;
            }
            double temp = b[i];
            b[i] = b[max_row];
            b[max_row] = temp;
        }

        // 检查主元是否为零
        if (fabs(A[i * (deg + 1) + i]) < 1e-10) {
            fprintf(stderr, "错误: 矩阵奇异，无法求解\n");
            free(A);
            free(b);
            return 1;
        }

        // 归一化当前行
        double pivot = A[i * (deg + 1) + i];
        for (int j = i; j <= deg; j++) {
            A[i * (deg + 1) + j] /= pivot;
        }
        b[i] /= pivot;

        // 消元
        for (int j = i + 1; j <= deg; j++) {
            double factor = A[j * (deg + 1) + i];
            for (int k = i; k <= deg; k++) {
                A[j * (deg + 1) + k] -= factor * A[i * (deg + 1) + k];
            }
            b[j] -= factor * b[i];
        }
    }

    // 回代求解
    for (int i = deg; i >= 0; i--) {
        coef[i] = b[i];
        for (int j = i + 1; j <= deg; j++) {
            coef[i] -= A[i * (deg + 1) + j] * coef[j];
        }
    }

    // 释放内存
    free(A);
    free(b);

    return 0;
}

/*
 * 多项式求值函数，等效于numpy.polyval
 * coef: 多项式系数数组
 * deg: 多项式阶数
 * x: 求值点
 * 返回值: 多项式在x点的值
 */
double polyval(const double* coef, int deg, double x) {
    double result = 0.0;
    for (int i = 0; i <= deg; i++) {
        result += coef[i] * pow(x, i);
    }
    return result;
}

// 示例使用
int test_polyfit_main() 
{
    // 示例数据
    double x[] = { 1.0, 2.0, 3.0, 4.0, 5.0 };
    double y[] = { 2.1, 3.9, 6.2, 8.1, 9.8 };
    int n = 5;  // 数据点数量
    int deg = 1;  // 多项式阶数（线性拟合）

    // 分配内存存储系数
    double* coef = (double*)malloc((deg + 1) * sizeof(double));
    if (!coef) {
        fprintf(stderr, "内存分配失败\n");
        return 1;
    }

    // 执行多项式拟合
    if (polyfit(x, y, n, deg, coef) != 0) {
        free(coef);
        return 1;
    }

    // 输出拟合结果
    printf("多项式系数 (从低次到高次): ");
    for (int i = 0; i <= deg; i++) {
        printf("%.6f ", coef[i]);
    }
    printf("\n");

    // 在x=6处求值
    double x_val = 3.0;
    double y_val = polyval(coef, deg, x_val);
    printf("在 x=%.1f 处的值: %.6f\n", x_val, y_val);

    // 释放内存
    free(coef);

    return 0;
}